# -*- coding: utf-8 -*-
# @Author: Ruban
# @License: Apache Licence
# @File: file_util.py

import os
import cv2
import numpy as np


def list_files(in_path):
    img_files = []
    mask_files = []
    gt_files = []
    for (dir_path, dir_names, file_names) in os.walk(in_path):
        for file in file_names:
            filename, ext = os.path.splitext(file)
            ext = str.lower(ext)
            if ext == '.jpg' or ext == '.jpeg' or ext == '.gif' or ext == '.png' or ext == '.pgm':
                img_files.append(os.path.join(dir_path, file))
            # elif ext == '.bmp':
            #     mask_files.append(os.path.join(dir_path, file))
            # elif ext == '.xml' or ext == '.gt' or ext == '.txt':
            #     gt_files.append(os.path.join(dir_path, file))
            # elif ext == '.zip':
            #     continue
    # img_files.sort()
    # mask_files.sort()
    # gt_files.sort()
    return img_files, mask_files, gt_files


def saveResult(img_file, img, boxes, dirname='./result/', verticals=None, texts=None):
    """ save text detection result one by one
    Args:
        img_file (str): image file name
        img (array): raw image context
        boxes (array): array of result file
            Shape: [num_detections, 4] for BB output / [num_detections, 4] for QUAD output
    Return:
        None
    """
    img = np.array(img)

    # make result file list
    filename, file_ext = os.path.splitext(os.path.basename(img_file))

    # result directory
    # res_file = dirname + "res_" + filename + '.txt'
    res_img_file = dirname + "res_" + filename + '.jpg'

    if not os.path.isdir(dirname):
        os.mkdir(dirname)

    # f = open(res_file, 'w')
    for i, box in enumerate(boxes):
        poly = np.array(box).astype(np.int32).reshape((-1))
        # strResult = ','.join([str(p) for p in poly]) + '\r\n'
        # f.write(strResult)

        poly = poly.reshape(-1, 2)
        cv2.polylines(img, [poly.reshape((-1, 1, 2))], True, color=(0, 0, 255), thickness=2)
        # ptColor = (0, 255, 255)
        # if verticals is not None:
        #     if verticals[i]:
        #         ptColor = (255, 0, 0)

        if texts is not None:
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            cv2.putText(img, "{}".format(texts[i]), (poly[0][0] + 1, poly[0][1] + 1), font, font_scale, (0, 0, 0),
                        thickness=1)
            cv2.putText(img, "{}".format(texts[i]), tuple(poly[0]), font, font_scale, (0, 255, 255), thickness=1)

    # Save result image
    cv2.imwrite(res_img_file, img)
